import asyncio
import cv2
import msgpack
import os
import ssl
import time
import numpy as np

from aioquic.asyncio.client import connect
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.events import StreamDataReceived, QuicEvent
from aioquic.asyncio.protocol import QuicConnectionProtocol

# --- 配置 ---
# 服务器地址：请确保这里与服务器证书的 CN (Common Name) 匹配。
# 如果服务器证书是为 'localhost' 生成的，这里也必须是 'localhost'。
HOST = "116.56.143.12" # <-- 请务必检查并确保您的文件是这一行！
PORT = 4433
# 服务器的自签名证书文件，客户端需要信任它
# 确保此文件与客户端脚本在同一目录下，且由服务器端生成
CERT_FILE = "server.pem"

# --- 客户端协议实现 ---
class VideoClientProtocol(QuicConnectionProtocol):
    """
    处理 QUIC 连接的客户端协议类。
    负责接收来自服务器的视频帧数据并进行显示。
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.frame_count = 0
        self.start_time = time.time()
        print("VideoClientProtocol 已初始化，准备接收视频流。")

    def quic_event_received(self, event: QuicEvent):
        """
        处理接收到的 QUIC 事件。
        当接收到新的流数据时，此方法会被调用。
        """
        # 检查事件是否为 StreamDataReceived (流数据接收事件)
        if isinstance(event, StreamDataReceived):
            try:
                # 使用 MessagePack 解包接收到的二进制数据
                # raw=False 表示将 MessagePack 的 bytes 类型数据转换为 Python 的 str 或 bytes
                unpacked_data = msgpack.unpackb(event.data, raw=False)
                frame_id = unpacked_data["frame_id"]
                frame_data = unpacked_data["data"]

                # 将接收到的图像字节数据转换为 NumPy 数组
                np_arr = np.frombuffer(frame_data, np.uint8)
                # 使用 OpenCV 从 NumPy 数组解码图像 (JPEG 格式)
                frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

                if frame is not None:
                    self.frame_count += 1
                    # 在窗口中显示接收到的视频帧
                    cv2.imshow("接收到的视频流 (按 'q' 退出)", frame)

                    # 每隔 30 帧打印一次当前帧率，提供实时反馈
                    if self.frame_count % 30 == 0:
                        elapsed_time = time.time() - self.start_time
                        fps = self.frame_count / elapsed_time if elapsed_time > 0 else 0
                        print(f"接收到帧 {frame_id}. 当前帧率: {fps:.2f} FPS")

                    # 检查用户是否按下了 'q' 键，如果按下则退出程序
                    # cv2.waitKey(1) 表示等待 1 毫秒，并处理 GUI 事件
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        print("用户请求退出。")
                        # 关闭 QUIC 连接
                        self.transport.close()
                        # 关闭所有 OpenCV 窗口
                        cv2.destroyAllWindows()
                        return
                else:
                    print(f"警告: 无法解码帧 {frame_id} 的图像数据。可能是数据损坏或格式错误。")

            except msgpack.exceptions.ExtraData as e:
                # 处理 MessagePack 解包时可能出现的额外数据错误
                print(f"MessagePack ExtraData 错误: {e}。数据可能不完整或格式错误。")
            except Exception as e:
                # 捕获并打印处理流数据时发生的任何其他错误
                print(f"处理流数据时发生错误: {e}")

        # 服务器为每个帧创建一个新的流并立即关闭，所以 StreamEnded 事件会频繁触发
        # 在这里可以根据需要添加处理逻辑，但对于视频流接收，通常可以忽略
        elif event.end_stream:
            pass

# --- 客户端运行逻辑 ---
async def run_client():
    """
    启动 QUIC 客户端并连接到服务器。
    """
    # 配置 QUIC 客户端参数
    configuration = QuicConfiguration(
        is_client=True,
        # ALPN (Application-Layer Protocol Negotiation) 协议，h3 是 HTTP/3 的标识
        # QUIC 客户端通常需要指定它
        alpn_protocols=["h3"],
        # 支持的 QUIC 协议版本
        supported_versions=["v1"]
    )

    # 创建 SSL 上下文以信任服务器的自签名证书
    # 在生产环境中，您应该使用由受信任的 CA 签发的证书，或更安全的证书验证方式。
    # 这里加载服务器生成的 server.pem 文件以进行测试。
    ssl_context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
    try:
        ssl_context.load_verify_locations(CERT_FILE)
    except FileNotFoundError:
        print(f"错误: 服务器证书 '{CERT_FILE}' 未找到。请确保它存在于当前目录。")
        print("如果您在测试环境中，并且确定要禁用证书验证 (不推荐用于生产环境)，可以取消注释以下两行：")
        # ssl_context.check_hostname = False # <-- 如果您确定要连接到 IP 地址，需要取消注释此行
        # ssl_context.verify_mode = ssl.CERT_NONE # <-- 如果您确定要连接到 IP 地址，需要取消注释此行
        return
    except Exception as e:
        print(f"加载 SSL 证书时发生错误: {e}")
        return

    print(f"正在连接到 QUIC 服务器 {HOST}:{PORT}...")
    try:
        # 使用 aioquic 的 connect 函数连接到服务器
        # create_protocol 参数指定了用于处理连接的协议类 (VideoClientProtocol)
        # local_port=0 让系统自动分配一个可用的本地端口
        async with connect(
            HOST,
            PORT,
            configuration=configuration,
            create_protocol=VideoClientProtocol,
            local_port=0
        ) as protocol:
            print(f"已成功连接到服务器 {HOST}:{PORT}")
            print("正在等待接收视频流...")
            # 客户端连接后，保持事件循环运行以持续接收服务器推送的流
            # 这里我们使用一个无限循环和 asyncio.sleep 来保持任务活跃
            # 直到连接被关闭 (例如，用户按下 'q' 或服务器断开连接)
            try:
                while True:
                    await asyncio.sleep(0.1) # 保持事件循环运行，处理接收到的帧
            except asyncio.CancelledError:
                print("客户端任务被取消。")
            finally:
                print("客户端连接已关闭。")
                cv2.destroyAllWindows() # 确保关闭所有 OpenCV 窗口

    except ConnectionRefusedError:
        print(f"错误: 连接被拒绝。请确保服务器正在 {HOST}:{PORT} 运行，并且网络可达。")
    except Exception as e:
        print(f"连接或运行客户端时发生错误: {e}")
        cv2.destroyAllWindows() # 确保在发生错误时也关闭窗口

# --- 主执行入口 ---
if __name__ == "__main__":
    # 确保 OpenCV 窗口在主线程中创建和管理，以避免 GUI 相关的警告或错误
    try:
        asyncio.run(run_client())
    except KeyboardInterrupt:
        print("客户端已通过 Ctrl+C 停止。")
    finally:
        # 确保在程序退出时关闭所有 OpenCV 窗口
        cv2.destroyAllWindows()
